#!/usr/bin/env python3
"""ROS 2 node that bridges Feetech servos <-> topic_based_ros2_control.

它现在包含一个完整的交互式校准程序。

它:
* 订阅 `/so101_follower/joint_commands` (sensor_msgs/JointState)
  并写入 Goal_Position 到舵机。
* 发布 `/so101_follower/joint_states` (使用校准过的零点)。
"""

from __future__ import annotations

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
import sys
import math
import yaml, pathlib
import json
from ament_index_python.packages import get_package_share_directory

try:
    import deepdiff  # type: ignore
except ImportError:  # create stub
    import types
    deepdiff_stub = types.ModuleType("deepdiff")
    class _DD(dict):
        pass
    deepdiff_stub.DeepDiff = _DD  # type: ignore
    sys.modules["deepdiff"] = deepdiff_stub

try:
    import tqdm  # type: ignore
except ImportError:  # create stub
    import types
    tqdm_stub = types.ModuleType("tqdm")
    def _tqdm(iterable=None, **kwargs):
        return iterable if iterable is not None else []
    tqdm_stub.tqdm = _tqdm  # type: ignore
    sys.modules["tqdm"] = tqdm_stub
# -----------------------------------------------------------------------------

from so101_hw_interface.motors.feetech.feetech import FeetechMotorsBus, OperatingMode
from so101_hw_interface.motors import Motor, MotorNormMode, MotorCalibration

PORT_DEFAULT = "/dev/ttyACM0"

JOINTS = {
    "1": {"id": 1, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
    "2": {"id": 2, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
    "3": {"id": 3, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
    "4": {"id": 4, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
    "5": {"id": 5, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
    "6": {"id": 6, "model": "sts3215", "mode": MotorNormMode.RANGE_0_100},
}

# INITIAL_RAW_POSITIONS = { ... }

CALIB_PATH = pathlib.Path.home() / ".calibrate" / "so101_follower_calibrate.json"

class MotorBridge(Node):
    def __init__(self):
        super().__init__("so101_motor_bridge")
        # Declare parameters so they can be overridden from launch/CLI
        self.declare_parameter("port", PORT_DEFAULT)
        self.declare_parameter("calib_file", str(CALIB_PATH)) 
        
        self.declare_parameter('calibrate', False)
        self.is_calibrate_mode = self.get_parameter('calibrate').get_parameter_value().bool_value
        if self.is_calibrate_mode:
            self.get_logger().warn(">>> 节点以 [校准模式] 启动 <<<")

        port = self.get_parameter("port").get_parameter_value().string_value
        if not port:
            port = PORT_DEFAULT

        # Build motor objects
        motors = {
            name: Motor(cfg["id"], cfg["model"], cfg["mode"])
            for name, cfg in JOINTS.items()
        }
        self.bus = FeetechMotorsBus(port, motors)
        self.joint_names = list(JOINTS.keys())

        self.get_logger().info(f"Connecting to Feetech bus on {port} …")
        try:
            self.bus.connect()
        except Exception as exc:  # noqa: BLE001
            self.get_logger().error(f"Could not open motor bus: {exc}")
            raise

        # Publishers / Subscribers
        self.joint_states_pub = self.create_publisher(JointState, "so101_follower/joint_states", 10)
        self.joint_commands_sub = self.create_subscription(
            JointState,
            "so101_follower/joint_commands",
            self._command_cb,
            10,
        )

        # State variables
        self.current_commands: dict[str, float] = {}
        self._home_offsets: dict[str, int] | None = None
        self._limits: dict[str, tuple[int, int]] | None = None
        self._is_read_turn = True # Flag to alternate between read and write
        self._steps_per_rad = 4096.0 / (2 * math.pi)

        self._reset_rad_offsets: dict[str, float] = {}
        
        # --- Type hint 现在是 MotorCalibration ---
        self.calibration_data: dict[str, MotorCalibration] | None = None 

        # --- 加载或运行校准 ---
        if not self._load_or_run_calibration():
            self.get_logger().error("校准/加载失败。正在关闭...")
            if self.bus:
                self.bus.disconnect()
            rclpy.shutdown()
            return

        # --- 从校准数据设置 home offsets 和 limits ---
        self._home_offsets = {name: cal.homing_offset for name, cal in self.calibration_data.items()}
        self._limits = {name: (cal.range_min, cal.range_max) for name, cal in self.calibration_data.items()}
        self.get_logger().info(f"Follower 臂已加载 home offsets: {self._home_offsets}")
        self.get_logger().info(f"Follower 臂已加载 limits: {self._limits}")
        
        # --- 现在运行 configure_motors 和 enable_torque ---
        self.get_logger().info("正在配置电机 (P/I/D, Torque)...")
        with self.bus.torque_disabled():
            self.bus.configure_motors() # 设置 Acceleration 等
            for motor in self.joint_names:
                self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
                self.bus.write("P_Coefficient", motor, 16)
                self.bus.write("I_Coefficient", motor, 0)
                self.bus.write("D_Coefficient", motor, 32)
                
                if motor == "6": # 假设 6 号是夹爪
                    self.get_logger().info("为夹爪 '6' 设置特殊扭矩限制。")
                    self.bus.write("Max_Torque_Limit", motor, 500)
                    self.bus.write("Protection_Current", motor, 250)
                    self.bus.write("Overload_Torque", motor, 25)
        
        self.bus.enable_torque()
        self.get_logger().info("电机总线已连接并配置完毕。")
        
        self.get_logger().info("正在设置初始复位位置...")
        reset_normalized_goals = {
            '1': -1.172161172161168,
            '2': -98.91304347826087,
            '3': 98.82724402345514,
            '4': -81.3107822410148,
            '5': -47.3015873015873,
            '6': 2.922077922077922
        }

        try:
            self.current_commands = reset_normalized_goals 
            
            # 写入初始位置
            self.bus.sync_write("Goal_Position", self.current_commands, normalize=True)
            
            self.get_logger().info(f"初始复位指令 (归一化值) 已设置。")
            self.get_logger().info(f"（保留两位小数）: { {k: round(v, 2) for k, v in self.current_commands.items()} }")

            # 2. (--- 新增 ---) 
            # 现在, 计算这些 归一化值 对应的 真实弧度值, 作为偏移量
            self.get_logger().info("正在计算复位位置的弧度偏移量...")
            self._reset_rad_offsets = {}
            for name, norm_val in reset_normalized_goals.items():
                motor = self.bus.motors[name]
                cal = self.calibration_data[name]
                min_ = cal.range_min
                max_ = cal.range_max
                home = cal.homing_offset

                # 2a. 归一化值 -> 原始步数(raw)
                raw = 0
                if (max_ - min_) == 0:
                    self.get_logger().error(f"电机 {name} 的校准范围无效 (min=max)。")
                    continue
                
                if motor.norm_mode == MotorNormMode.RANGE_M100_100:
                    raw = int(((norm_val + 100) / 200) * (max_ - min_) + min_)
                elif motor.norm_mode == MotorNormMode.RANGE_0_100:
                    raw = int((norm_val / 100) * (max_ - min_) + min_)
                
                # 2b. 原始步数(raw) -> 真实弧度(rad)
                rad = (raw - home) * (2 * math.pi) / 4096.0
                self._reset_rad_offsets[name] = rad
            
            self.get_logger().info(f"计算出的弧度偏移量 (保留两位): { {k: round(v, 2) for k, v in self._reset_rad_offsets.items()} }")

        except Exception as e:
            self.get_logger().error(f"设置初始复位位置失败: {e}")
            self.get_logger().warn("将使用 0.0 (归一化值) 作为默认初始位置。")
            self.current_commands = {n: 0.0 for n in self.joint_names}

        # Timer for periodic read/write (50 Hz)
        self.timer = self.create_timer(0.02, self._timer_cb)
        self.get_logger().info("Motor bridge 节点启动成功。")

    # ---------------------------------------------------------------------
    # 校准辅助函数
    # ---------------------------------------------------------------------

    def _save_calibration(self, calibration_data: dict[str, MotorCalibration]):
        """将完整的校准数据保存到 JSON 文件"""
        data_to_save = {
            name: {
                "id": cal.id,
                "drive_mode": cal.drive_mode,
                "homing_offset": cal.homing_offset,
                "range_min": cal.range_min,
                "range_max": cal.range_max,
            }
            for name, cal in calibration_data.items()
        }
        
        try:
            CALIB_PATH.parent.mkdir(parents=True, exist_ok=True)
            with open(CALIB_PATH, 'w') as f:
                json.dump(data_to_save, f, indent=4)
            self.get_logger().info(f"校准数据已保存到: {CALIB_PATH}")
        except Exception as e:
            self.get_logger().error(f"无法保存校准文件: {e}")

    def _run_calibration(self) -> dict[str, MotorCalibration] | None:
        """运行交互式校准程序以查找 Homing Offsets 和 Ranges"""
        self.get_logger().info("\n--- 开始新的 Follower 臂校准 ---")
        self.get_logger().info("!!! 注意: 此过程需要用户交互输入 !!!")
        try:
            self.bus.disable_torque()
            for motor_name in self.joint_names:
                self.bus.write("Operating_Mode", motor_name, OperatingMode.POSITION.value)
            self.get_logger().info("所有电机已设置为位置模式。")

            # 1. 捕获 Homing Offsets
            input(">>> (1/2) 请用手将 Follower 臂移动到其运动范围的【中间位置】, 然后按 ENTER ...")
            sys.stdout.flush() # 强制刷新
            homing_offsets = self.bus.set_half_turn_homings()
            self.get_logger().info(f"成功捕获 Homing Offsets: {homing_offsets}")

            # 2. 捕获 Ranges
            full_turn_motor = "5" 
            unknown_range_motors = [
                name for name in self.joint_names if name != full_turn_motor
            ]
            
            print(f"\n>>> (2/2) 现在, 请【依次】将【除 '{full_turn_motor}' 之外的所有】关节移动到它们的【最小值和最大值】。\n 完成后按 ENTER 停止记录...")

            range_mins, range_maxes = self.bus.record_ranges_of_motion(unknown_range_motors)

            range_mins[full_turn_motor] = 0
            range_maxes[full_turn_motor] = 4095
            
            self.get_logger().info(f"记录的相对 Range Mins: {range_mins}")
            self.get_logger().info(f"记录的相对 Range Maxes: {range_maxes}")

            full_calibration_data = {}
            for motor_name, m in self.bus.motors.items():
                full_calibration_data[motor_name] = MotorCalibration(
                    id=m.id,
                    drive_mode=0,
                    homing_offset=homing_offsets[motor_name],
                    range_min=range_mins[motor_name], 
                    range_max=range_maxes[motor_name], 
                )
            
            self.get_logger().info(f"转换后的 (相对) Raw Range Mins: { {k: v.range_min for k, v in full_calibration_data.items()} }")
            self.get_logger().info(f"转换后的 (相对) Raw Range Maxes: { {k: v.range_max for k, v in full_calibration_data.items()} }")

            # 5. 将校准数据写入电机固件
            self.get_logger().info("正在将完整校准数据写入电机固件...")
            self.bus.write_calibration(full_calibration_data) 
            self.get_logger().info("校准数据已写入电机。")

            # 6. 保存到本地文件
            self._save_calibration(full_calibration_data)
            
            return full_calibration_data
            
        except Exception as e:
            self.get_logger().error(f"校准过程中出错: {e}")
            import traceback
            self.get_logger().error(traceback.format_exc())
            return None

    def _load_or_run_calibration(self) -> bool:
        """
        - 如果 'calibrate' 参数为 True, 强制运行交互式校准。
        - 否则 (默认), 自动加载校准文件。
        - 如果文件不存在, 打印错误并失败。
        """
        
        # 模式 1: 强制校准 (用于手动运行)
        if self.is_calibrate_mode:
            self.get_logger().info("--- 运行交互式校准 (Follower) ---")
            self.calibration_data = self._run_calibration()
            if self.calibration_data is None:
                self.get_logger().error("交互式校准失败。")
                return False
            self.get_logger().info("交互式校准成功。")
            return True # 校准成功

        # 模式 2: 自动加载 (用于 Launch 文件)
        
        # 检查 calib_file 参数, 但优先使用 CALIB_PATH
        param_path_str = self.get_parameter("calib_file").get_parameter_value().string_value
        use_path = CALIB_PATH
        if param_path_str != str(CALIB_PATH) and ".yaml" in param_path_str:
             self.get_logger().warn(f"检测到参数指定了旧的 YAML 校准文件 ({param_path_str})。")
             self.get_logger().warn(f"将【忽略】此参数，并优先使用 JSON 校准：{CALIB_PATH}")
        
        if use_path.is_file():
            self.get_logger().info(f"找到校准文件: {use_path}。将自动加载此文件。")
            try:
                with open(use_path, 'r') as f:
                    loaded_data = json.load(f) # 这是一个 dict[str, dict]
                
                self.calibration_data = {}
                for name, data_dict in loaded_data.items():
                    if name in self.joint_names: 
                        self.calibration_data[name] = MotorCalibration(**data_dict) # 实例化
                
                if (not all(j in self.calibration_data for j in self.joint_names)):
                     self.get_logger().error("校准文件无效或不完整。请删除旧文件并运行手动校准。")
                     self.get_logger().error("  (在 OPi 上运行) ros2 run so101_hw_interface so101_motor_bridge --ros-args -p calibrate:=True")
                     return False
                else:
                    self.get_logger().info("已从文件加载校准数据。")
                    self.get_logger().info("正在将加载的校准数据写入电机固件...")
                    self.bus.write_calibration(self.calibration_data)
                    self.get_logger().info("校准数据已写入电机。")
                    return True
            except Exception as e:
                self.get_logger().error(f"加载/写入校准文件失败 ({e})。")
                self.get_logger().error("请删除损坏的文件并运行手动校准:")
                self.get_logger().error("  (在 OPi 上运行) ros2 run so101_hw_interface so101_motor_bridge --ros-args -p calibrate:=True")
                return False
        
        # 模式 2 (续): 文件未找到
        else:
            self.get_logger().error(f"未找到校准文件: {use_path}")
            self.get_logger().error("请在【单独的终端中】运行此节点以执行交互式校准:")
            self.get_logger().error("  (在 OPi 上运行) ros2 run so101_hw_interface so101_motor_bridge --ros-args -p calibrate:=True")
            self.get_logger().error("校准完成后, 重新启动 launch 文件。")
            return False


    # ---------------------------------------------------------------------
    # Callbacks
    # ---------------------------------------------------------------------
    def _command_cb(self, msg: JointState):
        """Store desired joint positions (converting rad -> normalized)."""
        try:
            for name, pos_rad in zip(msg.name, msg.position):
                if name in JOINTS:
                    
                    # 修复: 将 ros2_control 的相对指令 (pos_rad) 转换为 绝对弧度目标
                    # 0.0 (来自 controller) + 偏移量 = 我们的复位位置
                    true_goal_rad = pos_rad + self._reset_rad_offsets.get(name, 0.0)
                    
                    # --- 剩下的逻辑使用 true_goal_rad ---
                    motor = self.bus.motors[name]
                    cal = self.calibration_data[name]
                    home = cal.homing_offset
                    min_ = cal.range_min
                    max_ = cal.range_max

                    # 1. 将 绝对弧度(rad) 转换为 目标原始步数(raw)
                    target_raw = int((true_goal_rad * self._steps_per_rad) + home)
                    
                    # 2. 将 原始步数(raw) 限制在电机的物理范围内
                    target_raw = min(max(target_raw, min_), max_)

                    # 3. 将 目标原始步数(raw) 转换为 归一化值(normalized)
                    norm_val = 0.0
                    # 避免除以零
                    if (max_ - min_) == 0: 
                        self.get_logger().warn(f"电机 {name} 的校准范围无效 (min=max)。", throttle_duration_sec=5.0)
                        continue

                    if motor.norm_mode == MotorNormMode.RANGE_M100_100:
                        norm_val = ((target_raw - min_) / (max_ - min_)) * 200.0 - 100.0
                    elif motor.norm_mode == MotorNormMode.RANGE_0_100:
                        norm_val = ((target_raw - min_) / (max_ - min_)) * 100.0
                    else:
                        self.get_logger().warn(f"不支持的 norm_mode {motor.norm_mode} (在 _command_cb 中)", throttle_duration_sec=5.0)
                        continue
                    
                    # 4. 存储归一化值
                    self.current_commands[name] = norm_val
        
        except Exception as e:
             self.get_logger().error(f"转换 (弧度->归一化) 指令失败: {e}", throttle_duration_sec=5.0)

    def _timer_cb(self):
        # On each timer tick, we either do a read or a write, but never both.
        if self._is_read_turn:
            self._do_read()
        else:
            self._do_write()

        # Flip the flag for the next turn
        self._is_read_turn = not self._is_read_turn

    def _do_read(self):
        try:
            raw_positions = self.bus.sync_read("Present_Position", normalize=False)
        except Exception as exc:
            self.get_logger().warn(f"sync_read failed: {exc}")
            return

        # Publish current joint states
        js = JointState()
        js.header.stamp = self.get_clock().now().to_msg()
        js.name = list(JOINTS.keys())
        
        # 修复: 报告 (真实弧度 - 复位弧度偏移量)
        # 这样当机械臂在复位位置时, ros2_control 会看到 0.0
        js.position = [
            (raw - self._home_offsets.get(n, 0)) * (2 * math.pi) / 4096.0 - self._reset_rad_offsets.get(n, 0.0)
            for n, raw in raw_positions.items()
        ]
        self.joint_states_pub.publish(js)

    def _do_write(self):
        if self._home_offsets is None:
            self.get_logger().warn("Skipping write: home offsets not yet captured (校准未完成?)")
            return

        # On subsequent write cycles, send the latest commands from the topic.
        if not self.current_commands:
             # 第一次写入时, 读取当前位置, 并将其设为初始目标, 以防止突然移动。
            self.get_logger().info("第一次写入: 尚未收到外部指令。读取当前位置作为初始目标。", throttle_duration_sec=5.0)
            try:
                # 1. 新增一次读取, 获取当前电机位置 (相对步数)
                raw_positions = self.bus.sync_read("Present_Position")

                self.current_commands = raw_positions
                self.get_logger().info(f"当前位置: {raw_positions}")

            except Exception as exc:
                self.get_logger().warn(f"在 _do_write 中读取初始位置失败: {exc}")
                return

        try:
            raw_goals = self.current_commands
            # self.get_logger().info(f"Trying to write: {raw_goals}")
            self.bus.sync_write("Goal_Position", raw_goals, normalize=True)
        except Exception as exc:
            self.get_logger().warn(f"sync_write failed: {exc}")


def main():
    """Entry-point."""
    rclpy.init()
    node = None
    try:
        node = MotorBridge()
        # 只有在 __init__ 成功 (即校准成功) 且 bus 有效时才 spin
        if node.bus is not None and rclpy.ok() and node._home_offsets is not None:
            rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        if node and node.bus is not None:
            node.get_logger().info("Disconnecting from motor bus...")
            try:
                node.bus.disable_torque()
                node.get_logger().info("电机扭矩已禁用。")
            except Exception as e:
                node.get_logger().warn(f"禁用扭矩失败: {e}")
            node.bus.disconnect()
        if node and rclpy.ok():
            node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()